from sqlalchemy import create_engine
from sqlalchemy import text
import pandas as pd
import time


class MysqlEngine():
    def __init__(self, user_name, password, host, database='tmp', port=3307, schema=None, **kwargs):
        self.MAX_RETRY = 10
        url = f'mysql+mysqlconnector://{user_name}:{password}@{host}:{port}/{database}'
        if schema:
            self.engine = create_engine(url, connect_args={'options': f'-csearch_path={schema}'}, **kwargs)
        else:
            self.engine = create_engine(url, **kwargs)

    def run_sql(self, sql):
        with self.engine.begin() as conn:
            conn.execute(text(sql))

    def fetch_dat(self, sql):
        return pd.read_sql(sql, self.engine)

    def write_dat(self, dat, tbl_name, if_exists, index, schema=None):
        dat.to_sql(tbl_name, self.engine, if_exists=if_exists, index=index, schema=schema)

    def run_procedure(self, procedure_name):
        retry_times = 0
        while retry_times < self.MAX_RETRY:
            connection = self.engine.raw_connection()
            try:
                cursor = connection.cursor()
                cursor.callproc(procedure_name)
                cursor.close()
                connection.commit()
                break
            except BaseException as e:
                print(e)
                retry_times = retry_times + 1
                print(f'RUN PROCEDURE RETRYING {retry_times} TIMES...')
                time.sleep(3)
            finally:
                connection.close()
